
import math

from dexen.system import constants

class Selection(object):
    def __init__(self):
        pass
    

class EnvironmentalSelection(Selection):
    
    def __init__(self):
        pass

    def setObjectives(self, *attributes):
        self.objective_attributes = attributes

    def setSelectionRate(self, rate):
        self.selection_rate = rate
    
    def dominates(self, ind1, ind2):
        dominated = False
        for objective in self.objective_attributes:
            if objective[1] == constants.MINIMIZE:
                if ind2[objective[0]] < ind1[objective[0]]:
                    return False
                if ind1[objective[0]] < ind2[objective[0]]:
                    dominated = True

            if objective[1] == constants.MAXIMIZE:
                if ind2[objective[0]] > ind1[objective[0]]:
                    return False
                if ind1[objective[0]] > ind2[objective[0]]:
                    dominated = True
        if dominated:
            return True
        return False
    
    def _updatePareto(self, new_ind, non_dominated_ones):
        for ind in non_dominated_ones[:]:
            if self.dominates(new_ind, ind):
                non_dominated_ones.remove(ind)
            
    def filterDominatedOnes(self, individuals):
        assert len(self.objective_attributes) != 0
        non_dominated_ones = []
        non_dominated_ones.append(individuals[0])
        
        for ind1 in individuals[:]:
            is_dominated = False
            dominated = False
            for ind2 in non_dominated_ones:
                if self.dominates(ind2, ind1):
                    is_dominated = True
                    break
                if self.dominates(ind1, ind2):
                    dominated = True
                    break

            if dominated:
                self._updatePareto(ind1, non_dominated_ones)
                non_dominated_ones.append(ind1)
                
            if not dominated and not is_dominated:
                non_dominated_ones.append(ind1)
            
        
        while len(individuals) > 0:
            individuals.pop(0)
        
        individuals.extend(non_dominated_ones)    
        return non_dominated_ones

    def select(self, individuals):
        self.calculate_strength(individuals)
        self.calculate_raw_fitness(individuals)
        self.calculate_distances(individuals)
        self.calculate_density(individuals)
        self.calculate_fitness(individuals)
                    
        next_archive = [x for x in individuals if x['fitness'] < 1]
        dominated_ones = [x for x in individuals if x['fitness'] >= 1]
        archive_size = len(next_archive)
        N = int(len(individuals) * self.selection_rate)
        
        print "N is ", N
            
        if archive_size < N:
            dominated_ones.sort(key = lambda x:(x['fitness']))
            next_archive.extend(dominated_ones[:N - archive_size])
            #print "Archive is less than N"
        elif archive_size > N:
            for ind in individuals:
                ind['removed'] = False
            archive_truncation(next_archive, N)
            #print "Archive is greater than N"
    
        assert len(next_archive) == N
        self.pop_keys(individuals)
        return dominated_ones, next_archive
    
    def calculate_strength(self, individuals):
        for ind1 in individuals:
            ind1['strength'] = 0
            for ind2 in individuals:
                if self.dominates(ind1, ind2):
                    ind1['strength'] += 1
    
    def calculate_raw_fitness(self, individuals):
        for ind1 in individuals:
            ind1['raw_fitness'] = 0
            for ind2 in individuals:
                if self.dominates(ind2, ind1):
                    ind1['raw_fitness'] += ind2['strength']
    
    def calculate_fitness(self, individuals):
        for ind in individuals:
            ind['fitness'] = ind['raw_fitness'] + ind['density'] 
    
    def get_distance(self, ind1, ind2):
        result = 0.0
        for objective in self.objective_attributes:
            result += (ind1[objective[0]] - ind2[objective[0]])**2 
    
        result = math.sqrt(result)
        return result
        
    def calculate_distances(self, individuals):
        for ind1 in individuals:
            ind1['distances'] = []
            for ind2 in individuals:
                ind1['distances'].append({"ind" : ind2, "distance" : self.get_distance(ind1, ind2)})
            ind1['distances'].sort(key = lambda x:(x['distance']))
    
    def calculate_density(self, individuals):
        k = int(math.sqrt(len(individuals))) 
        for ind in individuals:
            ind['density'] = 1.0 / (ind['distances'][k]['distance'] + 1) #????? 1 or 2 
    
    def pop_keys(self, individuals):
        for ind in individuals:
            ind.pop('strength')
            ind.pop('raw_fitness')
            ind.pop('fitness')
            ind.pop('density')
            ind.pop('distances')
            if ind.has_key('removed'):
                ind.pop('removed')
    
    def get_distance_k(self, ind, k):
        cnt = 0
        pos = 0
        while cnt < k:
            assert pos < len(ind['distances'])
            if ind['distances'][pos]['ind']['removed'] == False and \
                            not(ind['distances'][pos]['ind'] is ind):
                cnt += 1
            pos += 1
        return ind['distances'][pos-1]['distance']
        
                
    def smaller(self, ind1, ind2, N):
        for k in range(N):
            dist1 = self.get_distance_k(ind1, k)
            dist2 = self.get_distance_k(ind2, k)
            if dist1 > dist2:
                return False
            if dist1 < dist2:
                return True
        return True
    
    def remove_individual(self, min_ind, individuals):
        for i, ind in enumerate(individuals):
            if ind is min_ind:
                individuals.pop(i)
                
    
    def archive_truncation(self, individuals, N):
        while (len(individuals) > N):
            min_ind = individuals[0]
            for ind in individuals[1:]:
                if self.smaller(ind, min_ind, len(individuals)):
                    min_ind = ind
            min_ind['removed'] = True
            #individuals.remove(min_ind)
            self.remove_individual(min_ind, individuals) 
